#include <stdio.h>
#include <stdlib.h>
#include <string.h>

/*
 * 9000000033
 * 对报文编码和解码
 * 哈夫曼树：权值小的放左子树，权值大的放右子树
 *	编码：左子树0，右子树1
 *	输入：
 *		字符集大小n (n <= 100)
 *		n个字符和n个权值（正整数），值越大则概率越大
 *		输入目标报文
 *	输出：
 *		编码后的二进制报文，占一行
 *		解码后的报文，占一行
 *		一个回车符
 *	例子：
 *	输入 

5 a b c d e 12 40 15 8 25
bbbaddeccbbb

 *	输出
00011111110111010110110000
bbbaddeccbbb

 *
 * */

#define DEBUG

#define MAX_STRING_SIZE 100
/* Pg147 of book, 6.12 */
typedef struct {
	unsigned int weight;
	unsigned int parent, lchild, rchild;
}HTNode, *HuffmanTree;

typedef char** HuffmanCode;

/*
在HT[1...i-1]选择parent为0且weight最小的两个节点
total: 当前一共有多少节点
其序号分别为s1, s2, HT[s1], HT[s2]
对于哈夫曼树，如果是不是叶子节点，
则一定同时包含左子树和右子树
*/

/*
	return index
*/
unsigned int find_min(HuffmanTree p, unsigned int total)
{
	unsigned int index, cur_weight;
	HuffmanTree cur;
	unsigned int ret;

	//find the first untreed weight
	//找到第一个还没有编进哈夫曼树的节点
	ret = 1;
	cur = &p[ret];
	cur_weight = cur->weight;
	for (index = 1; index <= total; index++, cur++) {
		if (!cur->parent) {
			ret = index;
			cur_weight = cur->weight;
			break;
		}
	}
	//寻找权重最小的节点
	for (; index <= total; index++, cur++) {
		if (cur->parent) //had been build in
			continue;
		if (cur_weight > cur->weight) {
			ret = index;
			cur_weight = cur->weight;
		}
	}

	return ret;
}

void Select(HuffmanTree p, unsigned int total, unsigned int* i1, unsigned int* i2)
{
	//有parent，代表这个节点已经编入huffman树
	*i1 = find_min(p, total);
	p[*i1].parent = 1; //为了找到下一个最小值，暂时设定它有parent
	*i2 = find_min(p, total);
	p[*i1].parent = 0; //restore
#ifdef DEBUG
	printf("select %d,%d\n", p[*i1].weight, p[*i2].weight);
#endif
}

//build huffman tree according weight array
void HuffmanCreate(HuffmanTree* pHT, unsigned int* w, unsigned int n)
{
	unsigned int i, m;
	HuffmanTree p, HT;
	unsigned int s1, s2;

	if (n <= 1)
		return;

	m = 2 * n - 1;
	//0号单元未用
	HT = (HuffmanTree)malloc((m + 1) * sizeof(HTNode));
	memset(HT, 0, (m + 1) * sizeof(HTNode));
	for (p = HT + 1, i = 1; i <= n; ++i, ++p, ++w)
		p->weight = *w;

	for (i = n + 1; i <= m; ++i) {
		//build huffman tree
		Select(HT, i - 1, &s1, &s2);
		HT[s1].parent = i;
		HT[s2].parent = i;
		HT[i].lchild = s1;
		HT[i].rchild = s2;
		HT[i].weight = HT[s1].weight + HT[s2].weight;
	}
	*pHT = HT;
}

/* encode
 * 存放n个字符的权值，
	* 构造哈夫曼树
	* 求出n个字符的哈夫曼编码HC
* */
void HuffmanCoding(HuffmanTree pHT, HuffmanCode *pHC, unsigned int *w, int n)
{
	int i, c, start;
	unsigned int f;
	HuffmanTree HT;
	HuffmanCode HC;
	char* cd;

	HT = pHT;
	HC = (HuffmanCode)malloc((n + 1) * sizeof(char*));
	cd = (char*)malloc(n * sizeof(char));
	memset(cd, 0, n * sizeof(char));
	//从叶子找到根，建立哈夫曼编码表，本质是一个二维数组
	//ary[n+1][n]，n个字符，编码最长为n-1个字符，加上末尾0，最长n
	//ary[0]未用
	//哈夫曼编码是一种变长度编码
	for (i = 1; i <= n; ++i) {
		start = n - 1;
		for (c = i, f = HT[i].parent; f != 0; c = f, f = HT[f].parent) {
			if (HT[f].lchild == c)
				cd[--start] = '0';
			else
				cd[--start] = '1';
		}
		HC[i] = (char*)malloc((n - start) * sizeof(char));
		strcpy(HC[i], &cd[start]);
	}

	free(cd);
	*pHC = HC;
}

void HuffmanDestroy(HuffmanTree HT, HuffmanCode HC, unsigned int n)
{
	unsigned int i;

	free(HT);

	for (i = 1; i <= n; ++i) {
		free(HC[i]);
	}
	free(HC);
}

//打印哈夫曼编码的函数
void PrintHuffmanCode(HuffmanCode htable, unsigned int* w, int n)
{
	printf("Huffman code : \n");
	for (int i = 1; i <= n; i++)
		printf("%d code = %s\n", w[i - 1], htable[i]);
}

char* encodeString(HuffmanCode htable, char* symbol, char* str, int n)
{
	unsigned int index;
	char* string;
#ifdef DEBUG
	int i = 0;
#endif
	unsigned int len = 0;

	char* resultStr = NULL;
	
	resultStr = malloc(1024 * 1024);
#ifdef DEBUG
	for (i = 1; i <= n; i++)
		printf("%c code = %s\n", symbol[i - 1], htable[i]);
	printf("\n");
	printf("original string is :\n%s\n\n", str);
	printf("encode result is :\n");
#endif

	string = str;
	while (*string) {
		index = *string - 'a' + 1;
#ifdef DEBUG
		printf("%s", htable[index]);
#endif
		sprintf(resultStr + len, "%s", htable[index]);
		len = len + strlen(htable[index]);
		string++;
	}
#ifdef DEBUG
	printf("\n\n");
#else
	printf("\n%s\n", resultStr);
#endif

	return resultStr;
}

char findSymbol(HuffmanCode htable, char* symbol_table, int n, char* string)
{
	int i;
	char c;

	c = symbol_table[0];
	for (i = 1; i <= n; i++) {
		if (!strcmp(htable[i], string)) {
			c = symbol_table[i - 1];
			break;
		}
	}

	return c;
}

/* 从根出发，按照字符 '0' 或'1'，
 * 找左孩子或者右孩子
 * 一直找到叶子，就是对应的字符
 * */
void decodeString(HuffmanTree pHT, HuffmanCode htable, char* symbol,
		char* encodestr, int n)
{
	unsigned int m, len;
	HuffmanTree p;
	char* str;
	char c;
	char decodechar[16]; //max 15 bit
	char originString[MAX_STRING_SIZE + 1];
	unsigned int oriStrLen = 0;

	if (!encodestr)
		return;

	m = n * 2 - 1;
	p = &pHT[m]; //p point to root
	str = encodestr;
	len = 0;
	while(*str){
		if(*str == '0')
			p = &pHT[p->lchild];
		else
			p = &pHT[p->rchild];
		decodechar[len] = *str;
		len++;
		str++;
		if (!p->lchild && !p->rchild) {
			//leaf
			decodechar[len] = 0;
			c = findSymbol(htable, symbol, n, decodechar);
			originString[oriStrLen] = c;
			oriStrLen++;
			//prepare next character
			len = 0;
			p = &pHT[m]; //p point to root
		}
	}
	originString[oriStrLen] = 0;
#ifdef DEBUG
	printf("decode string is:\n");
#endif
	printf("%s\n", originString);

	free(encodestr);
}

unsigned int weight[] = {12, 40, 15, 8, 25};
char symbol_table[] = "abcde";
char origin_string[MAX_STRING_SIZE + 1] = "bbbaddeccbbb";

/*
5 a b c d e 12 40 15 8 25 
bbbaddeccbbb
*/
void get_user_input(unsigned int *n, char* symtable, 
	unsigned int *weight, char* str)
{
	unsigned int i;

	//character set number
	scanf("%d", n);
	//n characters
	for (i = 0; i < *n; i++) {
		scanf("%s", &symbol_table[i]);
	}
	//weight
	for (i = 0; i < *n; i++) {
		scanf("%d", &weight[i]);
	}
	//original string
	scanf("%s", str);
}

//unsigned int weight[] = {2, 8, 7, 6, 5};
//unsigned int weight[] = {5, 29, 7, 8, 14, 23, 3, 11};
int main(int argc, char** argv)
{
	HuffmanTree pHT;
	HuffmanCode htable;
	unsigned int n;
	char* str;

#ifdef DEBUG
	n = sizeof(weight) / sizeof(unsigned int);
#else
	get_user_input(&n, symbol_table, weight, origin_string);
#endif

	HuffmanCreate(&pHT, weight, n);
	HuffmanCoding(pHT, &htable, weight, n);
	//PrintHuffmanCode(htable, weight, n);
	str = encodeString(htable, symbol_table, origin_string, n);
	decodeString(pHT, htable, symbol_table, str, n);
	HuffmanDestroy(pHT, htable, n);
	return 0;
}
